#ifndef _YOLO_THREAD_POOL_H_
#define _YOLO_THREAD_POOL_H_

#include "threadpool.h"
#include "postprocess.h"
#include "preprocess.h"
#include "utils/logging.h"
#include "yolov8.h"

#include <opencv2/opencv.hpp>

class YoloThreadPool
{
public:
    YoloThreadPool(const std::string model_path, const int thread_num,
                   const std::string label_path);

    ~YoloThreadPool();

    void Init();
    void DeInint();
    void AddInferenceTask(std::shared_ptr<cv::Mat> src, Preprocess &preprocess);

    int get_model_id();

    std::shared_ptr<cv::Mat> GetImageResultFromQueue();

    int GetTaseksSize();

private:
    int thread_num_{1};
    std::string model_path_{"null"};
    std::string label_path_{"null"};
    uint32_t id{0};
    std::unique_ptr<ThreadPool> pool_;
    std::queue<std::shared_ptr<cv::Mat>> image_results_;
    std::vector<std::shared_ptr<Yolov8Custom>> models_;
    std::mutex id_mutex_;
    std::mutex image_results_mutex_;
};

#endif